#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@author:hengk
@contact: hengk@foxmail.com
@datetime:2019-11-02 11:23
"""
from utils import load_config
from backbones import BackBoneFactory
from datasets.chartsets import ChartSets,ChartsCollate

import os
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable

def test():
    net.eval()
    evalutation = torch.zeros((1,10)).cuda()
    for batch_index, (images, labels) in enumerate(test_loader):
        images = Variable(images).cuda()
        labels = Variable(labels).cuda().long()
        outputs = net(images)
        outputs = torch.sigmoid(outputs)
        outputs[outputs>0.5] = 1
        outputs[outputs<=0.5] = 0
        ret = (outputs.long() == labels)
        evalutation = evalutation + torch.sum(ret , dim=0)
    print(evalutation/len(test_loader.dataset))
if __name__ == '__main__':

    cfg = load_config("config/config.yaml")
    test_paramters = cfg.test

    net = BackBoneFactory.create(test_paramters.backbone, test_paramters.pretrain, test_paramters.num_class)
    if (test_paramters.gpu != ""):
        if(test_paramters.gpu!=""):
            os.environ["CUDA_VISIBLE_DEVICES"] = test_paramters.gpu
        net.cuda()

    if(test_paramters.model_path!=""):
        checkpoint = torch.load(test_paramters.model_path)
        net.load_state_dict(checkpoint)

    collote = ChartsCollate()

    test_set = ChartSets(test_paramters.test_images, test_paramters.test_labels, test_paramters.long_size)
    test_loader = DataLoader(
        test_set,
        batch_size=test_paramters.batchsize,
        collate_fn=collote,
        shuffle=True,
        num_workers=3,
        drop_last=True,
        pin_memory=True)
    test()
